#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
import argparse
from scapy.all import Ether
import sys
sys.path.append(
    "/home/p4/tutorials/utils"
)

from time import sleep
import logging
import ptf
import ptf.testutils as tu
from ptf.base_tests import BaseTest
import p4runtime_sh.shell as sh
import p4runtime_shell_utils as shu
import scapy.all as scapy


logger = logging.getLogger(None)
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
# create formatter and add it to the handlers
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
# sys.path.append(
#     os.path.join(os.path.dirname(os.path.abspath(__file__)),
#                  '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import enum
from p4runtime_sh.context import P4Type
SWITCH_TO_HOST_PORT = 1
SWITCH_TO_SWITCH_PORT = 2

class PuntReason(enum.Enum):
    UNRECOGNIZED_OPCODE = 1
    OPERATION_RESPONSE  = 2
    ACTIVE_REQUEST = 3

class ControllerOptcode(enum.Enum):
    NO_OP          = 0
    READ_REGISTER  = 1
    WRITE_REGISTER = 2
    GENERATE_CONGA = 3


def int_to_hex(num):
    hex_str=hex(num)
    if len(hex_str[2:]) < 8:
        return '0x'+'.'.join(['0' for i in range(8-len(hex_str[2:]))])+hex_str[2:]
    
    return hex_str


def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.

    :param p4info_helper: the P4Info helper
    :param sw: the switch connection
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def printCounter(p4info_helper, sw, counter_name, index):
    """
    Reads the specified counter at the specified index from the switch. In our
    program, the index is the tunnel ID. If the index is 0, it will return all
    values from the counter.

    :param p4info_helper: the P4Info helper
    :param sw:  the switch connection
    :param counter_name: the name of the counter from the P4 program
    :param index: the counter index (in our case, the tunnel ID)
    """
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            print("%s %s %d: %d packets (%d bytes)" % (
                sw.name, counter_name, index,
                counter.data.packet_count, counter.data.byte_count
            ))

def printGrpcError(e):
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print("(%s)" % status_code.name, end=' ')
    traceback = sys.exc_info()[2]
    print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))


def get_exp_num_packetins(pktin, exp_num_packets, timeout_sec):
    """
        exp_num_packets
        timeout_spec
        @return: pktlist: 在pktin上接收到的报文
    """
    pktlist = []
    pktin.sniff(lambda p: pktlist.append(p), timeout=timeout_sec)
    # assert len(pktlist) == exp_num_packets
    return pktlist

        

# 1. 为交换机s1和s2以及s3实现一个controller
# 2. 控制器在收到消息需要产生conga报文的时候，生成conga报文返回给数据层
class Controller():
    CPU_PORT = 510

    def __init__(self,p4info_txt_fname,p4prog_binary_fname):
        """
        完成初始化工作，读取后缀为`.json`以及`.p4info.txtpb`格式的文件
        """
        # 判断文件是否存在

        # 初始化
        logging.info("start controller init.")
        self.p4info_txt_fname=p4info_txt_fname
        self.p4prog_binary_fname=p4prog_binary_fname
        
        p4info_data = shu.read_p4info_txt_file(p4info_txt_fname=p4info_txt_fname)
        
        # Create Python dicts from name to integer values, and integer
        # values to names, for the P4_16 serializable enum types
        # PuntReason_t and ControllerOpcode_t once here during setup.
        self.punt_reason_name2int, self.punt_reason_int2name = \
            shu.serializable_enum_dict(p4info_data, 'PuntReason_t')
        self.opcode_name2int, self.opcode_int2name = \
            shu.serializable_enum_dict(p4info_data, 'ControllerOpcode_t')

        self.p4info_obj_map = shu.make_p4info_obj_map(p4info_data)
        self.cpm_packetin_id2data = \
            shu.controller_packet_metadata_dict_key_id(self.p4info_obj_map,
                                                       "packet_in")

    def addController(self,
                      device_id:int,
                      grpc_addr:str,
                      verbose:bool=False):
        
        sh.setup(device_id=device_id,
                 grpc_addr=grpc_addr,
                 election_id=(1,0),
                 config=sh.FwdPipeConfig(self.p4info_txt_fname, self.p4prog_binary_fname),
                 verbose=verbose)

        # 连接建立完毕之后开始进行监听,todo: 使用多线程进行优化
        packetin=sh.PacketIn()
        
        while True:
            pktlist=get_exp_num_packetins(packetin,1,0.01) # 每隔10ms监听一次，大部分的数据中心网络回路收敛大概100ms左右
            # 解析收到的报文，根据收到的报文采取相应的动作
            if len(pktlist)>=1:
                pktdata=shu.decode_packet_in_metadata(self.cpm_packetin_id2data,
                                                      pktlist[0].packet)
                # 经过上面处理得到的pktdata为一个字典类型，key构成的集合为 (metadata, payload)
                # 其中metadata也是一个字典，里面又包含属性 (input_port, punt_reason, opcode, operand0, operand1, operand2, operand3)
                # analyse_pktdata(pktdata['metadata'],sh.PacketOut())
                pkt_metadata=pktdata['metadata']
                if(pkt_metadata['punt_reason']==PuntReason.ACTIVE_REQUEST.value and \
                        pkt_metadata['opcode']==ControllerOptcode.GENERATE_CONGA.value):
                    # 生成默认的conga报文，后续可以对路径进行优化 (todo)
                    
                    for i in [3,4]:
                        pktout=sh.PacketOut()
                        pktout.metadata['srcAddr']= int_to_hex(pkt_metadata['operand0'])
                        pktout.metadata['dstAddr']= int_to_hex(pkt_metadata['operand1'])
                        pktout.metadata['direction']='0'
                        pktout.metadata['LBTag']='%d' % i
                        pktout.metadata['CE']='0'
                        pktout.send()
            else:
                logging.debug("not receive packet")
        
# if __name__ == '__main__':
#     p4info_txt_fname='./build/ld_ml.p4.p4info.txtpb'
#     p4prog_binary_fname='./build/ld_ml.json'
#     controller=Controller(p4info_txt_fname,p4prog_binary_fname)
#     controller.addController(device_id=0,
#                              grpc_addr='localhost:50051',
#                              verbose=False)
